// Tencent is pleased to support the open source community by making
// 蓝鲸智云 - 监控平台/日志平台 (BlueKing - Monitor/Log) available.
// Copyright (C) 2017-2022 THL A29 Limited, a Tencent company. All rights reserved.
// Licensed under the MIT License (the "License"); you may not use this file except in compliance with the License.
// You may obtain a copy of the License at http://opensource.org/licenses/MIT
// Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
// an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.
//

package tokenchecker

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"math"
	"math/rand"
	"strconv"
	"strings"
	"sync"
	"unicode/utf8"

	"github.com/TencentBlueKing/bk-collector/define"
)

const (
	DecoderTypeRandom = "random"
	DecoderTypeFixed  = "fixed"
	DecoderTypeAcs256 = "aes256"
)

func SwitchTokenDecoder(c Config) TokenDecoder {
	switch c.Type {
	case DecoderTypeRandom:
		return RandomTokenDecoder(c)
	case DecoderTypeFixed:
		return FixedTokenDecoder(c)
	case DecoderTypeAcs256:
		return Aes256TokenDecoder(c)
	}
	return RandomTokenDecoder(c)
}

type TokenDecoder interface {
	Type() string
	Skip() bool
	Decode(s string) (define.Token, error)
}

// RandomTokenDecoder 根据配置生成随机 Token 用于测试场景

func RandomTokenDecoder(c Config) TokenDecoder {
	maxv := c.MaxVal
	if maxv > math.MaxInt32 || maxv <= 0 {
		maxv = math.MaxInt32
	}

	return randomTokenDecoder{
		token: define.Token{
			TracesDataId:  rand.Int31n(int32(maxv)),
			MetricsDataId: rand.Int31n(int32(maxv)),
			LogsDataId:    rand.Int31n(int32(maxv)),
		},
	}
}

type randomTokenDecoder struct {
	token define.Token
}

func (d randomTokenDecoder) Type() string                        { return DecoderTypeRandom }
func (d randomTokenDecoder) Skip() bool                          { return true }
func (d randomTokenDecoder) Decode(string) (define.Token, error) { return d.token, nil }

// FixedTokenDecoder 根据配置生成固定 Token 用于测试场景

func FixedTokenDecoder(c Config) TokenDecoder {
	return fixedTokenDecoder{
		token: define.Token{
			Original:      c.FixedToken,
			TracesDataId:  c.TracesDataId,
			MetricsDataId: c.MetricsDataId,
			LogsDataId:    c.LogsDataId,
			BizId:         c.BizId,
			AppName:       c.AppName,
		},
	}
}

type fixedTokenDecoder struct {
	token define.Token
}

func (d fixedTokenDecoder) Type() string                        { return DecoderTypeFixed }
func (d fixedTokenDecoder) Skip() bool                          { return true }
func (d fixedTokenDecoder) Decode(string) (define.Token, error) { return d.token, nil }

// Aes256TokenDecoder 使用 aes256 加盐算法 所有字段均由配置项指定

func Aes256TokenDecoder(c Config) TokenDecoder {
	return newAes256TokenDecoder(c)
}

type tokenWithError struct {
	Token define.Token
	Err   error
}

func newAes256TokenDecoder(c Config) *aes256TokenDecoder {
	h := sha256.New()
	h.Write([]byte(c.DecodedKey))
	key := h.Sum(nil)
	return &aes256TokenDecoder{
		salt:  c.Salt,
		key:   key,
		iv:    []byte(c.DecodedIv),
		cache: map[string]tokenWithError{},
	}
}

type aes256TokenDecoder struct {
	salt string
	key  []byte
	iv   []byte

	mut   sync.Mutex
	cache map[string]tokenWithError
}

func (d *aes256TokenDecoder) Type() string { return DecoderTypeAcs256 }
func (d *aes256TokenDecoder) Skip() bool   { return false }

func (d *aes256TokenDecoder) Decode(s string) (define.Token, error) {
	d.mut.Lock()
	defer d.mut.Unlock()

	v, ok := d.cache[s]
	if ok {
		return v.Token, v.Err
	}

	token, err := d.decode(s)
	if err == nil {
		token.Original = s
	}
	d.cache[s] = tokenWithError{
		Token: token,
		Err:   err,
	}
	return token, err
}

func (d *aes256TokenDecoder) decode(s string) (define.Token, error) {
	var token define.Token
	enc, err := base64.StdEncoding.DecodeString(s)
	if err != nil {
		return token, err
	}
	if len(enc) < aes.BlockSize {
		return token, fmt.Errorf("invalid prefix-enc len: %d", len(enc))
	}

	block, err := aes.NewCipher(d.key)
	if err != nil {
		return token, err
	}

	enc = enc[aes.BlockSize:]
	stream := cipher.NewCBCDecrypter(block, d.iv)
	stream.CryptBlocks(enc, enc)
	if len(enc) < aes.BlockSize {
		return token, fmt.Errorf("invalid suffix-enc len: %d", len(enc))
	}

	decodedRune, _ := utf8.DecodeRune(enc[len(enc)-1:])
	decoded := string(enc[:len(enc)-int(decodedRune)])

	split := strings.Split(decoded, d.salt)
	if len(split) < 5 {
		return token, fmt.Errorf("invalid split len: %d, str: %s", len(split), decoded)
	}

	metricsDataId, err := strconv.Atoi(split[0])
	if err != nil {
		return token, fmt.Errorf("invalid metrics dataid: %s", split[0])
	}
	tracesDataId, err := strconv.Atoi(split[1])
	if err != nil {
		return token, fmt.Errorf("invalid traces dataid: %s", split[1])
	}
	logsDataId, err := strconv.Atoi(split[2])
	if err != nil {
		return token, fmt.Errorf("invalid logs dataid: %s", split[2])
	}
	bizId, err := strconv.Atoi(split[3])
	if err != nil {
		return token, fmt.Errorf("invalid bizid: %s", split[3])
	}
	appName := split[4]

	return define.Token{
		TracesDataId:  int32(tracesDataId),
		MetricsDataId: int32(metricsDataId),
		LogsDataId:    int32(logsDataId),
		BizId:         int32(bizId),
		AppName:       appName,
	}, nil
}
